import threading
class LoadBalancer:
    def __init__(self):
        self.consumers = []
        self.pending_counts = {}  # 消费者ID到待处理任务数的映射
        self.lock = threading.Lock()

    def register_consumer(self, consumer):
        """注册新的消费者"""
        try:
            with self.lock:
                self.consumers.append(consumer)
                self.pending_counts[consumer.id] = 0
                # print(f"Consumer {consumer.id} 已注册")
        except Exception as e:
            print(f"错误: 注册消费者失败, 错误信息: {e}")

    def send_message(self, message):
        """分发消息到最优消费者"""
        try:
            with self.lock:
                if not self.consumers:
                    print("错误：没有可用的消费者")
                    return

                # 选择当前负载最小的消费者
                min_id = min(self.pending_counts, key=lambda k: self.pending_counts[k])
                selected_consumer = next(c for c in self.consumers if c.id == min_id)
                
                # 更新待处理计数
                self.pending_counts[min_id] += 1
                # print(f"消息 '{message}' 已发送给消费者 {min_id}，待处理数：{self.pending_counts[min_id]}")
                
                # 将消息放入消费者队列
                selected_consumer.queue.put(message)
        except KeyError as e:
            print(f"错误: 选择消费者失败, 错误信息: {e}")
        except Exception as e:
            print(f"错误: 分发消息失败, 错误信息: {e}")
    def notify_task_completed(self, consumer_id):
        """消费者完成处理时通知负载均衡器"""
        try:
            with self.lock:
                self.pending_counts[consumer_id] -= 1
                #print(f"消费者 {consumer_id} 完成处理，剩余待处理数：{self.pending_counts[consumer_id]}")
        except KeyError as e:
            print(f"错误: 消费者ID {consumer_id} 不存在, 错误信息: {e}")
        except Exception as e:
            print(f"错误: 通知任务完成失败, 错误信息: {e}")
balancer = LoadBalancer()